class EncoderCNN(nn.Module):
'''
The Encoder class
Args:
embed_size (int)
train_cnn (bool) if true trains the complete network
'''
def __init__(self, embed_size, train_cnn=False):
super(EncoderCNN, self).__init__()
self.train_cnn = train_cnn
self.cnn_model = models.resnet18(pretrained=True)
self.cnn_model.fc = nn.Linear(self.cnn_model.fc.in_features, embed_size) # resize outout shape
self.relu = nn.ReLU()
self.bn = nn.BatchNorm1d(embed_size)
def forward(self, images):
features = self.cnn_model(images)
# specify if the complete network should be trained or only the last one
for name, param in self.cnn_model.named_parameters():
if "fc.weight" in name or "fc.bias" in name:
param.requires_grad = True
else:
param.requires_grad = self.train_cnn
return self.bn(self.relu(features))
class DecoderRNN(nn.Module):
'''
The decoder class
Args:
embed_size (int): size of the embeddings
hidden_size (int): equal to embed size
vocab_size (int): equals the number of words in the vocab
num_layers (int): number of layers in the lstm
dropout (float): dropout ratio for the embeddings
pretrained_emb (bool): if true uses the pretrained glove embedding vectors
'''
def __init__(self, embed_size, hidden_size, vocab_size, num_layers, dropout=0, pretrained_emb=False):
super(DecoderRNN, self).__init__()
if pretrained_emb:
self.embedding = nn.Embedding.from_pretrained(vocab.vectors, freeze=True)
else:
self.embedding = nn.Embedding(vocab_size, embed_size) # ????????????????
self.lstm = nn.LSTM(input_size=embed_size,
hidden_size=hidden_size,
num_layers=num_layers,
# batch_first=True
)
self.linear = nn.Linear(hidden_size, vocab_size)
self.dropout = nn.Dropout(dropout)
def forward(self, features, captions, lengths):
# print("captions shape", np.shape(captions))
# embedding of captions
embeddings = self.dropout(self.embedding(captions))
# print("Embedded captions shape", np.shape(embeddings))
# print("features shape", np.shape(features))
embeddings = torch.cat((features.unsqueeze(1), embeddings), dim=1)
packed = pack_padded_sequence(embeddings, lengths, enforce_sorted=False, batch_first=True)
output_packed, hidden = self.lstm(packed) #(1, batch size, len_embedding)
output_padded, output_lengths = pad_packed_sequence(output_packed, batch_first=True)
outputs = self.linear(output_padded)
return outputs
class CNNtoRNN(nn.Module):
def __init__(self, embed_size, hidden_size, vocab_size, num_layers, dropout=0, pretrained_emb=False):
super(CNNtoRNN, self).__init__()
self.encoderCNN = EncoderCNN(embed_size=embed_size)
self.decoderRNN = DecoderRNN(embed_size=embed_size, hidden_size=hidden_size, vocab_size=vocab_size, num_layers=num_layers, pretrained_emb=pretrained_emb)
def forward(self, images, capitons, lengths):
features = self.encoderCNN(images)
outputs = self.decoderRNN(features, capitons, lengths)
return outputs
def caption_images(self, image, vocabulary, max_length = 30):
'''Creates a caption for a single image'''
caption_result = []
with torch.no_grad():
x = self.encoderCNN(image).unsqueeze(0)
states = None
for _ in range(max_length):
hidden, states = self.decoderRNN.lstm(x, states)
output = self.decoderRNN.linear(hidden.squeeze(0))
predicted = output.argmax(1)
caption_result.append(predicted.item())
x = self.decoderRNN.embedding(predicted).unsqueeze(0)
if vocabulary.itos[predicted.item()] == "<eos>":
break
return "".join([vocabulary.itos[idx] + " " for idx in caption_result])
def caption_image(model, path=image_path, filename=None, transform=image_transform):
'''
Creates a caption for a single image
Args:
model (CNNtoRNN): model class
path (str): path to image folder
filename (str): name of image, if None a random image is selected
transform (torchvision.transforms.Compose): the transformations of the image before feeding into the model
'''
model.eval()
if not filename:
filename = random.choice(os.listdir(path))
image = mpimg.imread(path + filename)
if transform is not None:
image_tensor = transform(image)
with torch.no_grad():
image_tensor = (image_tensor[None, ...]).to(device)
return filename, image, model.caption_images(image_tensor, vocab)
def func_save_model(model, path, name):
'''
Saves the model as state dict.
Args:
model (CNNtoRNN): the model class
path (str): path where to save the model
name of the document
'''
filename = "{}.pt".format(name)
print(filename)
if not os.path.exists(path):
os.makedirs(path)
torch.save(model, path + filename)
def func_save_log(log_dict, name):
'''
Saves the logging dictionary as pickle object
Args:
log_dict (dict): contains the logged metricts from the training
name (str): the name of the pickle file
'''
with open('{}.pickle'.format(name), 'wb') as f:
pickle.dump(log_dict, f, protocol=pickle.HIGHEST_PROTOCOL)
def func_open_log(name):
'''
Opens the logging dictionary from a pickle object
Args:
name (str): the name of the pickle file
Returns:
log_dict (dict): contains the logged metricts from the training
'''
with open('{}.pickle'.format(name), 'rb') as f:
return pickle.load(f)
